package com.slipper.system.ai.controller;

import com.slipper.ai.model.embedding.EmbeddingModelFactory;
import com.slipper.core.gson.GsonUtil;
import com.slipper.system.ai.bo.AiModelBO;
import com.slipper.system.ai.bo.KnowledgeBO;
import com.slipper.system.ai.bo.KnowledgeEmbeddingBO;
import com.slipper.system.ai.service.AiModelService;
import com.slipper.system.ai.service.KnowledgeEmbeddingService;
import com.slipper.system.ai.service.KnowledgeService;
import com.slipper.util.AffixUtil;
import jakarta.annotation.Resource;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.core.io.ByteArrayResource;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import com.slipper.system.ai.bo.KnowledgeFileBO;
import com.slipper.system.ai.service.KnowledgeFileService;
import com.slipper.constants.Constants;
import com.slipper.util.ResultUtil;
import com.slipper.util.Tool;

import java.util.ArrayList;
import java.util.List;

@RestController
@RequestMapping("/system/knowledgeFile")
public class KnowledgeFileController {
    @Resource
    private AiModelService aiModelService;
    @Resource
    private KnowledgeService knowledgeService;
    @Resource
    private KnowledgeFileService knowledgeFileService;
    @Resource
    private KnowledgeEmbeddingService knowledgeEmbeddingService;
    
    @PostMapping(value = "/queryBy")
    public String queryBy(@RequestBody KnowledgeFileBO bo) {
        if(Tool.isBlank(bo.getId())) {
            return ResultUtil.failureResult("数据获取失败！");
        }
        return ResultUtil.successResult(knowledgeFileService.queryBy(bo.getId()));
    }

    @PostMapping(value = "/queryList")
    public String queryList(@RequestBody KnowledgeFileBO bo) {
        return ResultUtil.successResult(knowledgeFileService.queryList(bo));
    }
    
    @PostMapping(value = "/queryPage")
    public String queryPage(@RequestBody KnowledgeFileBO bo) {
        return ResultUtil.successResult(knowledgeFileService.queryPage(bo));
    }

    @PostMapping(value = "/save")
    public String save(@RequestBody KnowledgeFileBO bo) {
        bo.setOperateId(Constants.getCurrentUser().getEmployeeId());
        return ResultUtil.saveResult(knowledgeFileService.save(bo), bo);
    }

    @PostMapping(value = "/delete")
    public String delete(@RequestBody KnowledgeFileBO bo) {
        return ResultUtil.deleteResult(knowledgeFileService.delete(bo));
    }

    @PostMapping(value = "/upload")
    public String upload(KnowledgeFileBO bo) throws Exception {
        if (bo.getFile() == null) {
            return ResultUtil.failureResult("请先选择要上传的文档！");
        }
        bo.setName(bo.getFile().getOriginalFilename());
        bo.setFileId(AffixUtil.saveFile(bo.getName(), bo.getFile().getInputStream()));
        bo.setStatus(Constants.NO);
        bo.setOperateId(Constants.getCurrentUser().getEmployeeId());
        bo.setFile(null);
        return ResultUtil.saveResult(knowledgeFileService.save(bo), bo);
    }

    @PostMapping(value = "/embedding")
    public String embedding(@RequestBody KnowledgeFileBO bo) throws Exception {
        if (Tool.isBlank(bo.getId())) {
            return ResultUtil.failureResult("参数不正确！");
        }
        KnowledgeFileBO knowledgeFileBO = knowledgeFileService.queryBy(bo.getId());
        if (knowledgeFileBO == null || Tool.isBlank(knowledgeFileBO.getFileId())) {
            return ResultUtil.failureResult("文件不存在！");
        }
        KnowledgeBO knowledgeBO = knowledgeService.queryBy(knowledgeFileBO.getKnowledgeId());
        if (knowledgeBO == null) {
            return ResultUtil.failureResult("知识库不存在！");
        }
        AiModelBO aiModelBO = aiModelService.queryBy(knowledgeBO.getModelId());
        if (aiModelBO == null) {
            return ResultUtil.failureResult("AI模型不存在！");
        }

        byte[] bytes = AffixUtil.readFile(knowledgeFileBO.getFileId());
        TikaDocumentReader reader = new TikaDocumentReader(new ByteArrayResource(bytes));
        List<Document> documentList = TokenTextSplitter.builder()
                .withChunkSize(knowledgeBO.getChunkSize())
                .withMinChunkSizeChars(knowledgeBO.getChunkMinSize())
                .build()
                .apply(reader.read());
        EmbeddingModel embeddingModel = EmbeddingModelFactory.create(aiModelBO);

        List<KnowledgeEmbeddingBO> embeddingList = new ArrayList<>();
        for (int i = 0; i < documentList.size(); i++) {
            Document document = documentList.get(i);
            KnowledgeEmbeddingBO embeddingBO = new KnowledgeEmbeddingBO();
            embeddingBO.setId(Tool.getUUID());
            embeddingBO.setKnowledgeId(knowledgeBO.getId());
            embeddingBO.setKnowledgeFileId(knowledgeFileBO.getId());
            embeddingBO.setOrdinal(i + 1);
            embeddingBO.setContent(document.getText());
            embeddingBO.setMetadata(GsonUtil.toJson(document.getMetadata()));
            embeddingBO.setEmbedding(GsonUtil.toJson(embeddingModel.embed(document)));
            embeddingBO.setOperateId(Constants.getCurrentUser().getEmployeeId());
            embeddingBO.setUpdateTime(System.currentTimeMillis());
            embeddingBO.setCreateTime(System.currentTimeMillis());
            embeddingList.add(embeddingBO);
        }
        knowledgeEmbeddingService.saveList(knowledgeFileBO.getId(), embeddingList);

        knowledgeFileBO.setStatus(Constants.YES);
        knowledgeFileService.save(knowledgeFileBO);
        return ResultUtil.successResult("向量化成功。", knowledgeFileBO);
    }
}
